#include "neural-net.h"
#include "assert.h"

#include <cstdio>
#include <map>
using std::map;

#include <iostream>
#include <sstream>
using namespace std;


static void GetNodeIndex(NeuralNetwork* n, const Node* node, NeuralNetwork::NodeType& nodetype, int& index) {
	if (find(n->inputs_.begin(), n->inputs_.end(), node) != n->inputs_.end()) {
		nodetype = NeuralNetwork::INPUT;
		index = find(n->inputs_.begin(), n->inputs_.end(), node) - n->inputs_.begin() + 1;
	}
	else if (find(n->hidden_nodes_.begin(), n->hidden_nodes_.end(), node) != n->hidden_nodes_.end()) {
		nodetype = NeuralNetwork::HIDDEN;
		index = find(n->hidden_nodes_.begin(), n->hidden_nodes_.end(), node) - n->hidden_nodes_.begin() + n->inputs_.size() + 1;
	}
	else if (find(n->outputs_.begin(), n->outputs_.end(), node) != n->outputs_.end()) {
		nodetype = NeuralNetwork::OUTPUT;
		index = find(n->outputs_.begin(), n->outputs_.end(), node) - n->outputs_.begin() + n->inputs_.size() + n->hidden_nodes_.size() + 1;
	}
}

static double ComputeRawValue(NeuralNetwork* n, const Node* node, int printFlag, NeuralNetwork::NodeType nodetype, int index) {
  // Take the weighted average of the inputs_.
  double total_weight = 0;

  stringstream input(stringstream::in | stringstream::out);
	stringstream weights(stringstream::in | stringstream::out);
	stringstream equation(stringstream::in | stringstream::out);
	
  for (size_t i = 0; i < node->inputs_.size(); ++i) {
		if (printFlag) {
			NeuralNetwork::NodeType nt;
			int temp_index;
			GetNodeIndex(n, node->inputs_[i], nt, temp_index);
		
			stringstream node_s;
			switch (nt) {
				case NeuralNetwork::INPUT:
					node_s << "Input " << temp_index;
		      break;
		    case NeuralNetwork::HIDDEN:
						node_s << "Hidden " << temp_index;
		      break;
		    case NeuralNetwork::OUTPUT:
						node_s << "Output " << temp_index;
		      break;
			}
			if (i != 0)
				input << ", ";
			input << node_s.str();
			stringstream weight_s;
			weight_s << "Weight_{" << index << temp_index << "}";
		
			weights << ", ";
			weights << weight_s.str();
			if (i != 0)
				equation << " + ";
			equation << node->inputs_[i]->transformed_value_ << " (" << node_s.str() << ")"
							<< " * "
							<< node->weights_[i]->value << " (" << weight_s.str() << ")";
		}
    total_weight += node->weights_[i]->value * node->inputs_[i]->transformed_value_;
  }
  total_weight += node->fixed_weight_->value;

	if (printFlag)
	cout << input.str()
		<< weights.str() <<endl
		<<"\t" << equation.str() <<endl;
  return total_weight;
}

static void UpdateInputWeights(NeuralNetwork* n, const Node& node, double delta, double learning_rate, int printFlag) {
	stringstream equation;
	int ni;
	if (printFlag) {
		NeuralNetwork::NodeType nt;
		GetNodeIndex(n, &node, nt, ni);
	
		stringstream node_s;
		switch (nt) {
			case NeuralNetwork::INPUT:
				node_s << "Input " << ni;
	      break;
	    case NeuralNetwork::HIDDEN:
					node_s << "Hidden " << ni;
	      break;
	    case NeuralNetwork::OUTPUT:
					node_s << "Output " << ni;
	      break;
		}
		
		cout << "\t" << "Updating weights going into " << node_s.str() << ":" <<endl;
	}
	
  for (size_t i = 0; i < node.inputs_.size(); ++i) {
    Weight* weight = node.weights_[i];
		double before = weight->value;
    weight->value += node.inputs_[i]->transformed_value_ * delta * learning_rate;

		if (printFlag) {
			NeuralNetwork::NodeType nt;
			int temp_index;
			GetNodeIndex(n, node.inputs_[i], nt, temp_index);
			
			cout << "\t\t" << "Weight_{" << ni << temp_index <<"}:"<< endl;
			cout << "\t\t\t" << learning_rate << " (alpha) * " << node.inputs_[i]->transformed_value_ << " (a_" << temp_index << ") * "
				<< delta << " (delta_" << ni << ") = " << node.inputs_[i]->transformed_value_ * delta * learning_rate  << endl;
			cout << "\t\t\t" << "Adding " << node.inputs_[i]->transformed_value_ * delta * learning_rate 
				<< " to " << before << " (Weight_{" << ni << temp_index <<"}) -> " << weight->value << endl;
		}
  }
  node.fixed_weight_->value += 1 * delta * learning_rate;
}

void NeuralNetwork::FeedForward(const Input& input, int printFlag) {
  CheckComplete();
  
	if (printFlag)
  	cout<<"FeedForward:"<<endl;
  
  // Instantiate all the input values.
 
  // cout << input.values.size() << " " <<inputs_.size() <<endl;
  assert(input.values.size() == inputs_.size());
  for (size_t i = 0; i < input.values.size(); ++i) {
    inputs_[i]->transformed_value_ = input.values[i];

		if (printFlag)
    	cout << " Input " << i+1 << ": " << input.values[i];
  }
	if (printFlag)
  cout << endl;

  // Now we need to compute the values for all nodes other than the inputs_.
  for (size_t i = 0; i < hidden_nodes_.size(); ++i) {
		int hi = i + inputs_.size() +1;
		if (printFlag)
			cout << "Feeding forward to Node " << hi << " (hidden) based on parents ";
		
    // Compute the value for this node.
    hidden_nodes_[i]->raw_value_ = ComputeRawValue(this, hidden_nodes_[i], printFlag, HIDDEN, hi);
    hidden_nodes_[i]->transformed_value_ = Sigmoid(hidden_nodes_[i]->raw_value_);

		if (printFlag)
			cout << "\t" << "Setting Node " << hi << " to " << hidden_nodes_[i]->raw_value_ <<endl
					<< "\t" << "Applying Sigmoid function to input to obtain output: g(" << hidden_nodes_[i]->raw_value_ << ") = "
					<<  hidden_nodes_[i]->transformed_value_ <<endl
					<< "\t" << "Setting output of Node " << hi << " to " << hidden_nodes_[i]->transformed_value_ <<endl;
  }
  for (size_t i = 0; i < outputs_.size(); ++i) {
		int oi = i+inputs_.size()+hidden_nodes_.size()+1;
		if (printFlag)
			cout << "Feeding forward to Node " << oi << " (output) based on parents ";
		
    // Compute the value for this node.
    outputs_[i]->raw_value_ = ComputeRawValue(this, outputs_[i], printFlag, OUTPUT, oi);
    outputs_[i]->transformed_value_ = Sigmoid(outputs_[i]->raw_value_);

		if (printFlag)
			cout << "\t" << "Setting Node " << oi << " to " << outputs_[i]->raw_value_ <<endl
					<< "\t" << "Applying Sigmoid function to input to obtain output: g(" << outputs_[i]->raw_value_ << ") = "
					<<  outputs_[i]->transformed_value_ <<endl
					<< "\t" << "Setting output of Node " << oi << " to " << outputs_[i]->transformed_value_ <<endl;
  }
}

void NeuralNetwork::Train(const vector<Input>& inputs,
                          const vector<Target>& targets,
                          double learning_rate,
                          int epochs,
													int printFlag) {
  CheckComplete();
  for (int i = 0; i < epochs; ++i) {
    for (size_t j = 0; j < inputs.size(); ++j) {
      Backprop(inputs[j], targets[j], learning_rate, printFlag);
    }
  }
}

void NeuralNetwork::Backprop(const Input& input,
                             const Target& target,
                             double learning_rate,
														 int printFlag) {
  CheckComplete();
  FeedForward(input, printFlag);
  map<Node*, double> node_to_delta;

	
	if (printFlag)
		cout<<"BackProp:"<<endl << "Learning rate: " << learning_rate <<endl;
	

  assert(outputs_.size() == target.values.size());
  // First we get all the deltas together, and then we update the weights.
  // Handle outputs_ first
	if (printFlag)
		cout << "Updating deltas for output nodes:" <<endl;
  for (size_t i = 0; i < outputs_.size(); ++i) {
    Node* node = outputs_[i];
    double delta =
      (target.values[i] - node->transformed_value_) * SigmoidPrime(node->raw_value_);
    node_to_delta[node] = delta;

		if (printFlag) {
			NeuralNetwork::NodeType nt;
			int temp_index;
			GetNodeIndex(this, node, nt, temp_index);
		
			stringstream node_s;
			switch (nt) {
				case NeuralNetwork::INPUT:
					node_s << "Input " << temp_index;
		      break;
		    case NeuralNetwork::HIDDEN:
						node_s << "Hidden " << temp_index;
		      break;
		    case NeuralNetwork::OUTPUT:
						node_s << "Output " << temp_index;
		      break;
			}
			stringstream equation;
			equation << "(" << target.values[i] << " (Target " << temp_index << ") - " << node->transformed_value_ << " (Output " << temp_index << ")) * "
				<< SigmoidPrime(node->raw_value_) << " (Sigmoid'(" << node->raw_value_ << "))";
			cout << "\t" << "Updating delta of " << node_s.str() << ":" << endl;
			cout << "\t\t" << equation.str() <<endl;
			cout << "\t\t" << "Setting delta of " << node_s.str() << " to " << delta <<endl;
		}
  }

	if (printFlag)
		cout << "Updating deltas for hidden nodes:" <<endl;
		
  for (int i = hidden_nodes_.size() - 1; i >= 0; --i) {
    Node* node = hidden_nodes_[i];

		stringstream equation;
		int hi;
		stringstream node_s;
		
		if (printFlag) {
			NeuralNetwork::NodeType nt;
			GetNodeIndex(this, node, nt, hi);
			
			switch (nt) {
				case NeuralNetwork::INPUT:
					node_s << "Input " << hi;
		      break;
		    case NeuralNetwork::HIDDEN:
						node_s << "Hidden " << hi;
		      break;
		    case NeuralNetwork::OUTPUT:
						node_s << "Output " << hi;
		      break;
			}
		}
		
    double delta = 0.0;
    for (size_t j = 0; j < node->forward_neighbors_.size(); ++j) {
      assert(node_to_delta.find(node->forward_neighbors_[j]) != node_to_delta.end());
      delta += node->forward_weights_[j]->value * node_to_delta[node->forward_neighbors_[j]];

			if (printFlag) {
				NeuralNetwork::NodeType nt;
				int temp_index;
				GetNodeIndex(this, node->forward_neighbors_[j], nt, temp_index);
		
				stringstream node_s;
				switch (nt) {
					case NeuralNetwork::INPUT:
						node_s << "Input " << temp_index;
			      break;
			    case NeuralNetwork::HIDDEN:
							node_s << "Hidden " << temp_index;
			      break;
			    case NeuralNetwork::OUTPUT:
							node_s << "Output " << temp_index;
			      break;
				}
				if (j != 0)
					equation << " + ";
				equation << node->forward_weights_[j]->value << " (Weight_{" << hi << temp_index << "}) * " << node_to_delta[node->forward_neighbors_[j]] << " (Delta " << hi << ")";
			}
    }
    delta *= SigmoidPrime(node->raw_value_);
    node_to_delta[node] = delta;

		if (printFlag) {
			cout << "\t" << "Updating delta of " << node_s.str() << ":" <<endl;
			cout << "\t\t" << equation.str() <<endl;
			cout << "\t\t" << "Setting delta of " << node_s.str() << " to " << delta <<endl;
		}
  }
	if(printFlag)
		cout << "Updating weights for output nodes:" <<endl;
  // Now that we have all the deltas, we can go ahead and update all the weights.
  for (size_t i = 0; i < outputs_.size(); ++i) {
    Node* node = outputs_[i];
    UpdateInputWeights(this, *node, node_to_delta[node],  learning_rate, printFlag);
  }
	if(printFlag)
		cout << "Updating weights for hidden nodes:" <<endl;
  for (size_t i = 0; i < hidden_nodes_.size(); ++i) {
    Node* node = hidden_nodes_[i];
    UpdateInputWeights(this, *node, node_to_delta[node],  learning_rate, printFlag);
  }
}
